from datetime import datetime

from colorama import Fore, Style

NO_PRINT = False
LOG_FILE = None
LOG_FILE_WRITE_F = None

def set_no_print(flag):
    global NO_PRINT
    NO_PRINT = flag

def set_log_file(file_path, mode='a'):
    global LOG_FILE
    LOG_FILE = file_path
    global LOG_FILE_WRITE_F
    LOG_FILE_WRITE_F = open(LOG_FILE, mode=mode, encoding='utf-8', buffering=1)
    LOG_FILE_WRITE_F.write("============DRAGON-LOG-START=================\n")

def _dragon_print(string_temp, flush=True):
    if NO_PRINT:
        return
    print(string_temp, flush=flush)
    if LOG_FILE_WRITE_F is not None:
        LOG_FILE_WRITE_F.write(string_temp + '\n')
    

def print_common(*args):
    if NO_PRINT:
        return
    string_temp = " ".join([str(arg) for arg in args])
    time_str = datetime.now().strftime(f'%Y-%m-%d %H:%M:%S DRAGON-PRINT ')
    string_temp = time_str + ' ' + string_temp
    _dragon_print(string_temp, flush=True)
def print_info(*args):
    if NO_PRINT:
        return
    string_temp = " ".join([str(arg) for arg in args])
    time_str = datetime.now().strftime(f'%Y-%m-%d %H:%M:%S DRAGON-INFO ')
    string_temp = time_str + ' ' + string_temp
    _dragon_print(Fore.GREEN + string_temp+ Style.RESET_ALL, flush=True)
def print_error(*args):
    if NO_PRINT:
        return
    string_temp = " ".join([str(arg) for arg in args])
    time_str = datetime.now().strftime(f'%Y-%m-%d %H:%M:%S DRAGON-ERROR ')
    string_temp = time_str + ' ' + string_temp
    _dragon_print(Fore.RED + string_temp + Style.RESET_ALL, flush=True)
def print_warning(*args):
    if NO_PRINT:
        return
    string_temp = " ".join([str(arg) for arg in args])
    time_str = datetime.now().strftime(f'%Y-%m-%d %H:%M:%S DRAGON-WARNING ')
    string_temp = time_str + ' ' + string_temp
    _dragon_print(Fore.YELLOW + string_temp + Style.RESET_ALL, flush=True)


def print_list(data: list):
    print_info('_________print_list_start_______________')
    for item in data:
        print_info(item)
    print_info('_________print_list_end____total:%d' % len(data))

def print_dict(data: dict):
    print_info('_________print_dict_start_______________')
    for k, v in data.items():
        print_info(f'{k} :\t{v}')
    print_info('_________print_dict_end____total:%d' % len(data))


def print_checkpoint(checkpoint):
    import torch
    if not isinstance(checkpoint, dict):
        assert isinstance(checkpoint, str), 'checkpoint must be a dict or a str'
        checkpoint = torch.load(checkpoint, map_location='cpu')
    print_info('_________print_checkpoint_start_______________')
    for k, v in checkpoint.items():
        print_info(f'{k} :\t{v.shape}')
    print_info('_________print_checkpoint_end____total:%d' % len(checkpoint))

def hello_DRAGON():
    """"""
    print_info("我是DRAGON，你好！")


class LimitPrinter:
    def __init__(self):
        self.max = 30
        self.now = 0

    def print(self, *args):
        text = ' '.join([str(x) for x in args])
        if self.now < self.max:
            print_info("LIMIT_PRINT: ", text)
            self.now += 1

    def set_max(self, max_in):
        self.max = max_in

    def reset(self):
        self.now = 0


global_limit_printer = LimitPrinter()


def print_limit(*text):
    global global_limit_printer
    global_limit_printer.print(*text)


# ------------------个性化 logging 相关-------------end

